package com.future;

import java.util.Arrays;

/**
 * Description: 498. 对角线遍历
 *
 * @author weiruibai.vendor
 * Date: 2022/12/27 10:27
 */
public class Solution_498 {

    private static Solution_498 instance = new Solution_498();

    public static void main(String[] args) {
        int[][] mat = new int[][]{
                {1, 2, 3},
                {4, 5, 6},
                {7, 8, 9}
        };
        int[] diagonalOrder = instance.findDiagonalOrder(mat);
        System.out.println(Arrays.toString(diagonalOrder));
    }

    public int[] findDiagonalOrder(int[][] mat) {
        int N = mat.length;
        int M = mat[0].length;
        int length = N * M;
        int[] ans = new int[length];
        int index = 0;
        int curX = 0;
        int curY = 0;
        ans[index++] = mat[curX][curY];
        boolean directDown = false;// 斜向下
        while (index < length) {
            if (!directDown && (curX == 0 || curY == M - 1)) { // 向上，来到第一行或者最后一列
                if (curX == 0 && curY < M-1) {
                    curY++;
                } else if (curY == M - 1) {
                    curX++;
                }
                directDown = true;
            } else if (directDown && (curX == N - 1 || curY == 0)) { // 向下，来到最后一行或第一列
                if (curY == 0 && curX < N-1) {
                    // 来到第一列，但是还未到最后一行
                    curX++;
                } else if (curX == N - 1) {
                    // 来到最后一行
                    curY++;
                }
                directDown = false;
            } else {
                if (directDown) {
                    curX++;
                    curY--;
                } else {
                    curX--;
                    curY++;
                }
            }
            ans[index++] = mat[curX][curY];
        }
        return ans;
    }
}
